module test_cuda_02_cpu_const

  implicit none

  integer, parameter :: r8 = 8
  integer, parameter :: ns = 1
  integer, parameter :: nblk = 80
  integer, parameter :: nthr = 512
  integer, parameter :: nx = nblk * nthr
  integer, parameter :: nt = 1000000
  real(r8), parameter :: dx = 1.0d0 / nx
  real(r8), parameter :: c = 1.0d0
  real(r8), parameter :: dt = 0.8d0 * dx / c

end module test_cuda_02_cpu_const

module test_cuda_02_cpu_kernels

  use test_cuda_02_cpu_const

contains

  subroutine upwind(q, f)

    real(r8), intent(in) :: q(1-ns:nx+ns)
    real(r8), intent(inout) :: f(1-ns:nx+ns)

    integer i

    !$omp parallel do simd schedule(static)
    do i = 1, nx
      f(i) = 0.5d0 * (c * (q(i+1) + q(i)) - abs(c) * (q(i+1) - q(i)))
    end do

  end subroutine upwind

  subroutine update(f, q_old, q_new)

    real(r8), intent(in) :: f(1-ns:nx+ns)
    real(r8), intent(in) :: q_old(1-ns:nx+ns)
    real(r8), intent(inout) :: q_new(1-ns:nx+ns)

    integer i

    !$omp parallel do simd schedule(static)
    do i = 1, nx
      q_new(i) = q_old(i) - dt / dx * (f(i) - f(i-1))
    end do

  end subroutine update

  subroutine fill_halo(x)

    real(r8), intent(inout) :: x(1-ns:nx+ns)

    integer i

    do i = 1 - ns, 0
      x(i) = x(i+nx)
    end do
    do i = nx + 1, nx + ns
      x(i) = x(i-nx)
    end do

  end subroutine fill_halo

end module test_cuda_02_cpu_kernels

program test_cuda_02_cpu

  use test_cuda_02_cpu_const
  use test_cuda_02_cpu_kernels

  implicit none

  real(r8), allocatable :: x_h(:)
  real(r8), allocatable :: q_h(:)
  real(r8), allocatable :: q_d(:,:)
  real(r8), allocatable :: f_d(:)

  integer i, old, new

  allocate(x_h(1-ns:nx+ns))
  allocate(q_h(1-ns:nx+ns))
  allocate(q_d(1-ns:nx+ns,2))
  allocate(f_d(1-ns:nx+ns))

  do i = lbound(x_h, 1), ubound(x_h, 1)
    x_h(i) = (i - 1) * dx
    if (x_h(i) >= 0.05 .and. x_h(i) <= 0.3) then
      q_h(i) = 1
    else
      q_h(i) = 0
    end if
  end do

  q_d(:,1) = q_h

  old = 1; new = 2
  do i = 1, nt
    call upwind(q_d(:,old), f_d)
    call fill_halo(f_d)
    call update(f_d, q_d(:,old), q_d(:,new))
    call fill_halo(q_d(:,new))
    call time_advance()
  end do

  print *, sum(x_h(1:nx) * q_d(1:nx,old))

  deallocate(x_h)
  deallocate(q_h)
  deallocate(q_d)
  deallocate(f_d)

contains

  subroutine time_advance()

    integer tmp

    tmp = old
    old = new
    new = tmp

  end subroutine time_advance

end program test_cuda_02_cpu
